#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
File: content_download.py
Author: StellaJL
Date: 2025-09-16 15:45:34
Version：1.0.0
Description: 异步内容下载器
"""



import aiohttp
import asyncio
import os
import yaml
from logger import get_logger


class ContentDownloader:
    """
    异步内容下载器类，返回URL请求内容
    注意，使用时需要使用上下文管理器，如：
    - async with ContentDownloader(timeout=10, config_path='config.yaml') as downloader:
        - content = await downloader.download(url)
    - 会根据响应头的Content-Type返回不同类型的内容，如：
        - 图片：返回bytes
        - JSON：返回dict
        - XML/HTML：返回str
        - 其他：返回str
        - 其他情况都会返回None，如有需要，请自行修改返回内容/添加新的解析方式
    """
    
    def __init__(self, timeout: int = 10, config_path: str = 'config.yaml') -> None:
        """
        初始化下载器
        
        Args:
            timeout: 请求超时时间（秒），默认为10秒
            config_path: 配置文件路径，默认为'config.yaml'
        """
        self.config_path = config_path
        self.timeout: aiohttp.ClientTimeout = aiohttp.ClientTimeout(total=timeout)
        self.logger = get_logger('downloader')
        
        # 初始化配置
        self._load_config()
        
        # 初始化HTTP组件
        self._session: aiohttp.ClientSession = None
        self._connector: aiohttp.TCPConnector = None
        self._initialized = False
        
        # 请求间隔相关
        self._last_request_time = 0
        self._request_interval = 0
    
    def _load_config(self):
        """加载配置文件"""
        try:
            if not os.path.exists(self.config_path):
                raise FileNotFoundError(f"配置文件 {self.config_path} 未找到。")
            
            with open(self.config_path, 'r', encoding='utf-8') as f:
                self._config = yaml.safe_load(f)
            
            # 加载HTTP配置
            http_config = self._config.get('http', {})
            
            # 设置User-Agent
            self.headers = {
                'User-Agent': http_config.get('user_agent', 
                    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36')
            }
            
            # 设置连接器配置
            connector_config = http_config.get('connector', {})
            self._connector_limit = connector_config.get('limit', 100)
            self._connector_limit_per_host = connector_config.get('limit_per_host', 30)
            
            # 加载爬虫配置中的请求间隔
            crawler_config = self._config.get('crawler', {})
            self._request_interval = crawler_config.get('interval', 0)  # 默认无间隔
            
            self._initialized = True
            
        except Exception as e:
            self.logger.error(f"加载配置失败: {e}")
            raise

    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self._initialize_session()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.close()

    async def _initialize_session(self):
        """初始化HTTP会话和连接器"""
        if self._session is None or self._session.closed:
            # 创建TCP连接器
            self._connector = aiohttp.TCPConnector(
                limit=self._connector_limit,
                limit_per_host=self._connector_limit_per_host
            )
            
            # 创建ClientSession
            self._session = aiohttp.ClientSession(
                timeout=self.timeout,
                connector=self._connector,
                headers=self.headers
            )
            
            self.logger.info(f"HTTP会话初始化完成 - 连接限制: {self._connector_limit}, 每主机限制: {self._connector_limit_per_host}")

    async def _wait_for_interval(self):
        """等待请求间隔时间"""
        if self._request_interval <= 0:
            return
        
        current_time = asyncio.get_event_loop().time()
        time_since_last_request = current_time - self._last_request_time
        
        if time_since_last_request < self._request_interval:
            wait_time = self._request_interval - time_since_last_request
            self.logger.debug(f"等待请求间隔: {wait_time:.2f}秒")
            await asyncio.sleep(wait_time)
        
        self._last_request_time = asyncio.get_event_loop().time()

    async def close(self):
        """关闭HTTP会话和连接器"""
        if self._session and not self._session.closed:
            await self._session.close()
            self.logger.info("HTTP会话已关闭")
        
        if self._connector and not self._connector.closed:
            await self._connector.close()
            self.logger.info("TCP连接器已关闭")

    async def download(self, url: str) -> str | None:
        """
        异步下载给定URL的内容
        
        Args:
            url: 要下载的URL地址
            
        Returns:
            Optional[str]: 成功时返回内容字符串，失败时返回None
        """
        self.logger.debug(f"开始下载: {url}")
        
        # 等待请求间隔
        await self._wait_for_interval()
        
        # 确保会话已初始化
        if self._session is None or self._session.closed:
            await self._initialize_session()
        
        try:
            async with self._session.get(url) as response:
                response.raise_for_status()  # 如果状态码不是2xx，会抛出异常
                content_type = response.headers.get('Content-Type', '').lower()
                content_length = response.headers.get('Content-Length', 'unknown')
                
                self.logger.info(f"成功下载 {url}, 状态码: {response.status}, 内容类型: {content_type}, 大小: {content_length}")
                
                if 'image' in content_type:
                    content = await response.read()
                    self.logger.debug(f"下载图片数据: {len(content)} 字节")
                    return content
                elif 'json' in content_type:
                    content = await response.json()
                    self.logger.debug(f"下载JSON数据: {len(str(content))} 字符")
                    return content
                elif 'xml' in content_type or 'html' in content_type:
                    content = await response.text()
                    self.logger.debug(f"下载XML/HTML数据: {len(content)} 字符")
                    return content
                else:
                    # 其他情况直接返回文本
                    content = await response.text()
                    self.logger.debug(f"下载文本数据: {len(content)} 字符")
                    return content
                
        except aiohttp.ClientError as e:
            self.logger.error(f"下载 {url} 时发生客户端错误: {e}")
        except asyncio.TimeoutError:
            self.logger.warning(f"下载 {url} 超时")
        except Exception as e:
            self.logger.error(f"下载 {url} 时发生未知错误: {e}")
            
        return None
